import numpy as np
# import matplotlib.pyplot as plt
# from pathlib import Path
# from tqdm import tqdm
# import argparse
import base64
import json
# import sys
# import os
import os.path as osp
import io
import math
import imgviz
from PIL import Image, ImageDraw
import cv2, os
from pathlib import Path
# import yaml
# import pandas as pd

# image to array
def img_b64_to_arr(img_b64):
    f = io.BytesIO()
    f.write(base64.b64decode(img_b64))
    img_arr = np.array(Image.open(f))
    return img_arr

# Convert labels shape to array
def shapes_to_label(img_shape, shapes, label_name_to_value, type='class'):    
    assert type in ['class', 'instance']
    mask_shape_array = np.zeros(img_shape[:2], dtype=np.int32)
    
    if type == 'instance':
        ins = np.zeros(img_shape[:2], dtype=np.int32)
        instance_names = ['_background_']
    
    for shape in shapes:
        
        points = shape['points'] # mask points
        label = shape['label'] # mask label name -> String
        shape_type = shape.get('shape_type', None) # polygon or something else
        
        if type == 'class':
            cls_name = label # mask label name -> String
            
        elif type == 'instance':
            cls_name = label.split('-')[0]
            if label not in instance_names:
                instance_names.append(label)
            ins_id = instance_names.index(label)
        
        cls_name = cls_name.strip()
        class_id = label_name_to_value[cls_name] # mask label id in 0, 1, 2...
        mask = shape_to_mask(img_shape[:2], points, shape_type) # mask shape (H * W)
        mask_shape_array[mask] = class_id # set the area mask not zero to mask_id
        
        if type == 'instance':
            ins[mask] = ins_id

    if type == 'instance':
        return mask_shape_array, ins
    
    return mask_shape_array

# Convert shape to mask
def shape_to_mask(img_shape, points, shape_type=None, line_width=10, point_size=5):
    
    mask = np.zeros(img_shape[:2], dtype=np.uint8) # generate numpy (H * W)

    mask = Image.fromarray(mask) # convert nupmy to PIL

    draw = ImageDraw.Draw(mask)
    
    xy = [tuple(point) for point in points] # generate (x,y)
    
    if shape_type == 'circle':
        assert len(xy) == 2, 'Shape of shape_type=circle must have 2 points'
        (cx, cy), (px, py) = xy
        d = math.sqrt((cx - px) ** 2 + (cy - py) ** 2)
        draw.ellipse([cx - d, cy - d, cx + d, cy + d], outline=1, fill=1) 
    elif shape_type == 'rectangle':
        assert len(xy) == 2, 'Shape of shape_type=rectangle must have 2 points'
        draw.rectangle(xy, outline=1, fill=1)
    elif shape_type == 'line':
        assert len(xy) == 2, 'Shape of shape_type=line must have 2 points'
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'linestrip':
        draw.line(xy=xy, fill=1, width=line_width)
    elif shape_type == 'point':
        assert len(xy) == 1, 'Shape of shape_type=point must have 1 points'
        cx, cy = xy[0]
        r = point_size
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=1, fill=1)
    else:
        mask = np.zeros(img_shape[:2], dtype=np.uint8) # generate numpy (H * W)
        cv2.drawContours(mask, [np.array(points, dtype=np.int32)], 0, 255, -1)  # 在二值图上绘制轮廓
        mask = np.array(mask, dtype=bool)
        return mask
        # assert len(xy) > 2, 'Polygon must have points more than 2'
        
        # draw.polygon(xy=xy, outline=1, fill=1)
    
    mask = np.array(mask, dtype=bool)
  
    return mask


# Save mask to image
def save_mask_to_image(filename, mask_with_class_id):
    
    if osp.splitext(filename)[1] != '.png':
        filename += '.png'

    # Assume label ranses [-1, 254] for int32, and [0, 255] for uint8 as VOC.
    if mask_with_class_id.min() >= -1 and mask_with_class_id.max() < 255:
        mask_with_class_id_pil = Image.fromarray(mask_with_class_id.astype(np.uint8), mode='P')
        
        colormap = imgviz.label_colormap(n_label=128) # generate color map
        
        mask_with_class_id_pil.putpalette(colormap.flatten()) # save mask to image conposed with diffierent color
        mask_with_class_id_pil.save(filename)
        
    else:
        raise ValueError(
            '[%s] Cannot save the pixel-wise class label as PNG. '
            'Please consider using the .npy format.' % filename
        )

def json_to_image(img_shape, json_path, toImage=False, visual=False, mask_label='box'):
    """ Covert json labels to segementation masks
    Args:
        img_path: input img path
        json_path: json file full name
        toImage:
        visual: whether show the mask on origin image
        mask_label: label name for the mask
    Returns:
        save masks to given output folder
    Raises:
        IOError: An error occurred accessing ablumentation object.
    """
    
    with open(json_path, 'r+', encoding="utf-8") as f:
        """ Load json file and update dictionary order """
        data = json.load(f)

    data['shapes'] = sorted(data['shapes'], key = lambda x: x['label'], reverse=False)
    # imageData = data.get('imageData')
    # # data = json.load(open(os.path.join(path_in, json_file)), encoding="utf-8-sig")
    
    # """ Convert imageData to numpy array """
    # if not imageData:
    #     if Path(data['imagePath']).stem:
    #         img = np.asarray(PIL.Image.open(img_path))
    #         # with open(img_path, 'rb') as f:
    #         #     imageData = f.read()
    #         #     imageData = base64.b64encode(imageData).decode('utf-8')
    #         #     img = img_b64_to_arr(imageData)
    #     else:
    #         raise ValueError('Can not find image file for this mask file.')
    # else:
    #     img = img_b64_to_arr(imageData)
    # print(f"01 The dimension of origin image is {img.shape}")

    """ Generate labels in the type of label encoder """
    # label encode each mask class
    mask_to_label_encode = {'_background_': 0}
    for i in range(len(data['shapes'])):
        shape_slice = data['shapes'][i]
        label_slice = shape_slice['label']
        if label_slice not in mask_to_label_encode.keys():
            mask_to_label_encode[label_slice] = 1
    # mask_to_label_encode = {'_background_': 0, mask_label: 1} 
    # mask_to_label_encode = {'_background_': 0, '01_bonding': 1, 
    #                         '02_spot': 2, '03_spatter': 3, '04_slag': 4}
    
    
    #         for shape in sorted(data['shapes'], key = lambda x: x['label'], reverse=False): # data['shapes'] -> list, sort by label class
    #             mask_label_name = shape['label'] # Only one shape each round

    #             if mask_label_name in mask_to_label_encode:
    #                 mask_label_value = mask_to_label_encode[mask_label_name] # mask_label_name means 'Joint', 'Line', or something else
    #             else:
    #                 mask_label_value = len(mask_to_label_encode) # If mask not include in list, add new mask label value
    #                 mask_to_label_encode[mask_label_name] = mask_label_value

    """ Generate labels in the type of classID """        
    mask_with_class_id = shapes_to_label(img_shape, data['shapes'], mask_to_label_encode) # mask's value are class id
    
    return mask_with_class_id


def mask_to_json(mask, imagePath, label_name='1'):
    # if os.path.basename(mask_path) == "0016E5_07959.png":
    #     print('t')
    # gray = cv2.imread(mask_path, 0)
    # gray = cv2.cvtColor(png, cv2.COLOR_BGR2GRAY)
    # img_file_path = os.path.join(img_path, os.path.basename(file).split('.')[0] + '.jpg')
    # img = Image.open(img_file_path)
    # imgData = img_tobyte(img)
    # imagePath = Path(mask_path).name.with_suffix('.jpg')
    dic = {"version": "4.5.6", "flags": {}, "shapes": list(), "imagePath": imagePath, "imageData": 'null',
           "imageHeight": mask.shape[0], "imageWidth": mask.shape[1]}
        
    class_dict = {label_name: 1}

    for k, v in class_dict.items():

        binary = mask.copy().astype(np.uint8)
        # binary[binary != v] = 0
        binary[binary > 0] = 255
 
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
        for contour in contours:
            temp = list()
            if len(contour) < 4:
                continue
            for point in contour:
                temp.append([float(point[0][0]), float(point[0][1])])
            dic["shapes"].append({"label": k, "points": temp, "group_id": None,
                                  "shape_type": "polygon", "flags": {}})
 
    return dic
